"""
RISC-V代码生成器核心模块

主要的代码生成器类，整合所有功能模块，提供高级和底层的代码生成接口。
"""

from .registers import RegisterManager
from .instructions import InstructionSet
from .utils import LabelManager, CodeSection, StackManager


class RISCVCodeGenerator:
    """RISC-V汇编代码生成器主类"""
    
    def __init__(self):
        self.register_manager = RegisterManager()
        self.instructions = InstructionSet(self)
        self.label_manager = LabelManager()
        self.stack_manager = StackManager(self.register_manager)
        
        # 代码段管理
        self.sections = {
            'data': CodeSection('data'),
            'text': CodeSection('text'),
            'main': CodeSection('main')
        }
        self.current_section = self.sections['main']
        
        # 代码生成状态
        self.code_lines = []
        self.comments = True  # 是否生成注释
        
    def emit(self, instruction, comment=None):
        """输出一条汇编指令"""
        if comment and self.comments:
            line = f"    {instruction:<20} # {comment}"
        else:
            line = f"    {instruction}"
        
        self.code_lines.append(line)
        self.current_section.add_instruction(line)
    
    def emit_label(self, label):
        """输出标签"""
        line = f"{label}:"
        self.code_lines.append(line)
        self.current_section.add_label(label)
    
    def emit_comment(self, comment):
        """输出注释"""
        if self.comments:
            line = f"    # {comment}"
            self.code_lines.append(line)
            self.current_section.add_instruction(line)
    
    def set_section(self, section_name):
        """切换当前代码段"""
        if section_name in self.sections:
            self.current_section = self.sections[section_name]
        else:
            raise ValueError(f"Unknown section: {section_name}")
    
    def get_register(self, name_or_reg):
        """获取寄存器，支持字符串名称或Register对象"""
        if isinstance(name_or_reg, str):
            return self.register_manager.get_register(name_or_reg)
        return name_or_reg
    
    def allocate_register(self, preferred=None):
        """分配寄存器"""
        return self.register_manager.allocate_register(preferred)
    
    def free_register(self, register):
        """释放寄存器"""
        self.register_manager.free_register(register)
    
    # ========== 高级控制流接口 ==========
    
    def create_function(self, name, save_registers=True):
        """创建函数"""
        self.emit_label(name)
        if save_registers:
            self.stack_manager.setup_function_prologue(self)
        return FunctionContext(self, name, save_registers)
    
    def create_loop(self, condition_reg=None, limit_reg=None, counter_reg=None):
        """创建循环结构"""
        return LoopContext(self, condition_reg, limit_reg, counter_reg)
    
    def create_if(self, condition_reg, condition_type='ne', compare_reg=None):
        """创建条件判断结构"""
        return IfContext(self, condition_reg, condition_type, compare_reg)
    
    # ========== 向量操作接口 ==========
    
    def vector_add(self, dst_addr, src1_addr, src2_addr, length, element_size=4):
        """向量加法操作"""
        with VectorAddContext(self, dst_addr, src1_addr, src2_addr, length, element_size) as ctx:
            return ctx
    
    # ========== 数据定义接口 ==========
    
    def define_data_word(self, label, *values):
        """定义字数据"""
        current = self.current_section
        self.set_section('data')
        self.emit_label(label)
        for value in values:
            self.emit(f".word {value}")
        self.current_section = current
    
    def define_data_string(self, label, string):
        """定义字符串数据"""
        current = self.current_section
        self.set_section('data')
        self.emit_label(label)
        self.emit(f'.string "{string}"')
        self.current_section = current
    
    def define_data_space(self, label, size):
        """定义空间"""
        current = self.current_section
        self.set_section('data')
        self.emit_label(label)
        self.emit(f".space {size}")
        self.current_section = current
    
    # ========== 代码输出接口 ==========
    
    def generate_code(self):
        """生成完整的汇编代码"""
        code = []
        
        # 数据段
        if self.sections['data'].instructions:
            code.append(".section .data")
            code.extend(self.sections['data'].instructions)
            code.append("")
        
        # 代码段
        code.append(".section .text")
        code.append(".globl _start")
        code.append("")
        
        # 主程序
        if self.sections['main'].instructions:
            code.extend(self.sections['main'].instructions)
        
        # 其他代码段
        if self.sections['text'].instructions:
            code.extend(self.sections['text'].instructions)
        
        return '\n'.join(code)
    
    def save_to_file(self, filename):
        """保存代码到文件"""
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(self.generate_code())
    
    def clear(self):
        """清空所有生成的代码"""
        self.code_lines.clear()
        for section in self.sections.values():
            section.clear()
        self.register_manager.free_all_temp_registers()


class FunctionContext:
    """函数上下文管理器"""
    
    def __init__(self, code_gen, name, save_registers):
        self.code_gen = code_gen
        self.name = name
        self.save_registers = save_registers
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.save_registers:
            self.code_gen.stack_manager.setup_function_epilogue(self.code_gen)


class LoopContext:
    """循环上下文管理器"""
    
    def __init__(self, code_gen, condition_reg=None, limit_reg=None, counter_reg=None):
        self.code_gen = code_gen
        self.condition_reg = condition_reg
        self.limit_reg = limit_reg
        self.counter_reg = counter_reg or code_gen.allocate_register('t0')
        self.loop_start, self.loop_end = code_gen.label_manager.create_loop_labels()
    
    def __enter__(self):
        # 初始化计数器
        if self.counter_reg:
            self.code_gen.instructions.li(self.counter_reg, 0)
        
        # 循环开始标签
        self.code_gen.emit_label(self.loop_start)
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 递增计数器
        if self.counter_reg and self.limit_reg:
            self.code_gen.instructions.addi(self.counter_reg, self.counter_reg, 1)
            self.code_gen.instructions.blt(self.counter_reg, self.limit_reg, self.loop_start)
        elif self.condition_reg:
            self.code_gen.instructions.bne(self.condition_reg, 'zero', self.loop_start)
        
        # 循环结束标签
        self.code_gen.emit_label(self.loop_end)
        
        # 释放计数器寄存器
        if self.counter_reg:
            self.code_gen.free_register(self.counter_reg)
    
    def continue_loop(self):
        """跳转到循环开始"""
        self.code_gen.instructions.j(self.loop_start)
    
    def break_loop(self):
        """跳出循环"""
        self.code_gen.instructions.j(self.loop_end)


class IfContext:
    """条件判断上下文管理器"""
    
    def __init__(self, code_gen, condition_reg, condition_type='ne', compare_reg=None):
        self.code_gen = code_gen
        self.condition_reg = condition_reg
        self.condition_type = condition_type
        self.compare_reg = compare_reg or 'zero'
        self.if_else, self.if_end = code_gen.label_manager.create_if_labels()
        self.in_else = False
    
    def __enter__(self):
        # 根据条件类型生成分支指令
        if self.condition_type == 'eq':
            self.code_gen.instructions.bne(self.condition_reg, self.compare_reg, self.if_else)
        elif self.condition_type == 'ne':
            self.code_gen.instructions.beq(self.condition_reg, self.compare_reg, self.if_else)
        elif self.condition_type == 'lt':
            self.code_gen.instructions.bge(self.condition_reg, self.compare_reg, self.if_else)
        elif self.condition_type == 'ge':
            self.code_gen.instructions.blt(self.condition_reg, self.compare_reg, self.if_else)
        elif self.condition_type == 'gt':
            self.code_gen.instructions.ble(self.condition_reg, self.compare_reg, self.if_else)
        elif self.condition_type == 'le':
            self.code_gen.instructions.bgt(self.condition_reg, self.compare_reg, self.if_else)
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self.in_else:
            self.code_gen.emit_label(self.if_else)
        self.code_gen.emit_label(self.if_end)
    
    def else_block(self):
        """进入else块"""
        self.code_gen.instructions.j(self.if_end)
        self.code_gen.emit_label(self.if_else)
        self.in_else = True
        return self


class VectorAddContext:
    """向量加法上下文管理器"""
    
    def __init__(self, code_gen, dst_addr, src1_addr, src2_addr, length, element_size=4):
        self.code_gen = code_gen
        self.dst_addr = dst_addr
        self.src1_addr = src1_addr
        self.src2_addr = src2_addr
        self.length = length
        self.element_size = element_size
        
        # 分配寄存器
        self.dst_reg = code_gen.allocate_register('t0')
        self.src1_reg = code_gen.allocate_register('t1')
        self.src2_reg = code_gen.allocate_register('t2')
        self.counter_reg = code_gen.allocate_register('t3')
        self.length_reg = code_gen.allocate_register('t4')
        self.temp_reg = code_gen.allocate_register('t5')
    
    def __enter__(self):
        # 加载数组地址
        self.code_gen.instructions.la(self.dst_reg, self.dst_addr)
        self.code_gen.instructions.la(self.src1_reg, self.src1_addr)
        self.code_gen.instructions.la(self.src2_reg, self.src2_addr)
        
        # 设置长度和计数器
        if isinstance(self.length, int):
            self.code_gen.instructions.li(self.length_reg, self.length)
        else:
            self.code_gen.instructions.mv(self.length_reg, self.length)
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 释放寄存器
        registers_to_free = [
            self.dst_reg, self.src1_reg, self.src2_reg,
            self.counter_reg, self.length_reg, self.temp_reg
        ]
        for reg in registers_to_free:
            self.code_gen.free_register(reg)
    
    def generate_loop(self):
        """生成向量加法循环"""
        with self.code_gen.create_loop(counter_reg=self.counter_reg, limit_reg=self.length_reg):
            # 计算当前元素的偏移量
            offset_reg = self.code_gen.allocate_register('t6')
            self.code_gen.instructions.slli(offset_reg, self.counter_reg, 2)  # offset = counter * 4
            
            # 加载两个操作数
            val1_reg = self.code_gen.allocate_register('a0')
            val2_reg = self.code_gen.allocate_register('a1')
            
            # 计算地址并加载
            self.code_gen.instructions.add(self.temp_reg, self.src1_reg, offset_reg)
            self.code_gen.instructions.lw(val1_reg, 0, self.temp_reg)
            
            self.code_gen.instructions.add(self.temp_reg, self.src2_reg, offset_reg)
            self.code_gen.instructions.lw(val2_reg, 0, self.temp_reg)
            
            # 执行加法
            self.code_gen.instructions.add(val1_reg, val1_reg, val2_reg)
            
            # 存储结果
            self.code_gen.instructions.add(self.temp_reg, self.dst_reg, offset_reg)
            self.code_gen.instructions.sw(val1_reg, 0, self.temp_reg)
            
            # 释放临时寄存器
            self.code_gen.free_register(offset_reg)
            self.code_gen.free_register(val1_reg)
            self.code_gen.free_register(val2_reg) 